#include <opencv2/opencv.hpp>
#include <iostream>
#include <vector>
#include <string> 
#include "vibe3.h"
#include <random>

#define BLOSIZE 32
#define BLOCKSIZE 32
#define DEFAULT_RANDOM_SAMPLE 16
#define DEFAULT_NUM_SAMPLES 20
#define DEFAULT_MAXHITCOUNT 50
#define DEFAULT_MIN_MATCH 2

// 定义颜色数组
const std::vector<cv::Scalar> COLORS = {
    {56, 0, 255},
    {226, 255, 0},
    {0, 94, 255},
    {0, 37, 255},
    {0, 255, 94},
    {255, 226, 0},
    {0, 18, 255},
    {255, 151, 0},
    {170, 0, 255},
    {0, 255, 56},
    {255, 0, 75},
    {0, 75, 255},
    {0, 255, 169},
    {255, 0, 207},
    {75, 255, 0},
    {207, 0, 255},
    {37, 0, 255},
    {0, 207, 255},
    {94, 0, 255},
    {0, 255, 113},
    {255, 18, 0},
    {255, 0, 56},
    {18, 0, 255},
    {0, 255, 226},
    {170, 255, 0},
    {255, 0, 245},
    {151, 255, 0},
    {132, 255, 0},
    {75, 0, 255},
    {151, 0, 255},
    {0, 151, 255},
    {132, 0, 255},
    {0, 255, 245},
    {255, 132, 0},
    {226, 0, 255},
    {255, 37, 0},
    {207, 255, 0},
    {0, 255, 207},
    {94, 255, 0},
    {0, 226, 255},
    {56, 255, 0},
    {255, 94, 0},
    {255, 113, 0},
    {0, 132, 255},
    {255, 0, 132},
    {255, 170, 0},
    {255, 0, 188},
    {113, 255, 0},
    {245, 0, 255},
    {113, 0, 255},
    {255, 188, 0},
    {0, 113, 255},
    {255, 0, 0},
    {0, 56, 255},
    {255, 0, 113},
    {0, 255, 188},
    {255, 0, 94},
    {255, 0, 18},
    {18, 255, 0},
    {0, 255, 132},
    {0, 188, 255},
    {0, 245, 255},
    {0, 169, 255},
    {37, 255, 0},
    {255, 0, 151},
    {188, 0, 255},
    {0, 255, 37},
    {0, 255, 0},
    {255, 0, 170},
    {255, 0, 37},
    {255, 75, 0},
    {0, 0, 255},
    {255, 207, 0},
    {255, 0, 226},
    {255, 245, 0},
    {188, 255, 0},
    {0, 255, 18},
    {0, 255, 75},
    {0, 255, 151},
    {255, 56, 0},
    {245, 255, 0}
};


cv::Mat apply_mask_to_image(const cv::Mat& image, const cv::Mat& mask) {
    // 检查 mask 是否为单通道
    if (mask.channels() != 1) {
        throw std::invalid_argument("Mask must be a single channel image");
    }

    // 创建一个与输入图像相同形状的全黑图像
    cv::Mat masked_image = cv::Mat::zeros(image.size(), image.type());

    // 创建随机数生成器
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, COLORS.size() - 1);

    // 随机选择一个颜色
    int color_index = dis(gen);
    cv::Scalar color = COLORS[color_index];

    // 将mask区域设为随机选择的颜色
    masked_image.setTo(color, mask == 255);

    // 将原始图像的非mask区域保留
    image.copyTo(masked_image, mask == 0);

    return masked_image;
}

cv::Mat fill_empty(const cv::Mat& imgtmp, cv::Mat& SegModel) {
    // 查找轮廓
    std::vector<std::vector<cv::Point>> contours;
    std::vector<cv::Vec4i> hierarchy;
    cv::findContours(imgtmp, contours, hierarchy, cv::RETR_TREE, cv::CHAIN_APPROX_NONE);

    for (size_t i = 0; i < contours.size(); i++) {
        // 一级父轮廓
        int father = hierarchy[i][3];
        // 二级父轮廓
        int grandpa = father >= 0 ? hierarchy[father][3] : -1;

        // 有一级父轮廓，无两级父轮廓，说明该轮廓是等级为1的轮廓，即需要的前景空洞区域
        if (father >= 0 && grandpa == -1) {
            // 填充面积 <= 100 的前景空洞区域
            double area = cv::contourArea(contours[i]);
            if (area <= 100) {
                cv::drawContours(SegModel, contours, i, cv::Scalar(255), -1);
            }
        }

        // 无父轮廓，说明该轮廓是最外围轮廓，即等级为0的前景斑点区域
        if (father == -1) {
            // 填充面积 < 10 的前景斑点区域
            double area = cv::contourArea(contours[i]);
            if (area < 10) {
                cv::drawContours(SegModel, contours, i, cv::Scalar(0), -1);
            }
        }
    }
    return SegModel;
}

std::string GetBaseName(const std::string& path) {
    std::string::size_type pos = path.find_last_of("/\\");
    if (pos != std::string::npos) {
        return path.substr(pos + 1);
    } else {
        return path; // 如果没有找到斜杠，返回整个路径
    }
}





int main() 
{
    int H = 1080;
    int W = 1920;
    int num_samples = DEFAULT_NUM_SAMPLES;
    int random_sample = DEFAULT_RANDOM_SAMPLE;
    int max_hitcount = DEFAULT_MAXHITCOUNT;
    int min_match = DEFAULT_MIN_MATCH;
    std::string RTSP;

    float* d_gray;
    float* d_samples;
    int* d_img_count;
    int* d_target_hit_count;
    int* h_mask_img;
    int* d_mask_img;
    // RTSP = "/home/xk/cv/vibeplus/BackgroundSplit-OpenCV-master/Video/Camera Road 01.avi";
    RTSP = "/home/xk/cv/2.mp4";
    // RTSP = "/data/汉丹线数据-原始数据/智洋/20241029/崩塌落石/20241029_091630_落石_80cm80cm_白天_晴_k411+530.mp4";
    // RTSP = "/data/汉丹线数据-原始数据/智洋/20241029/崩塌落石/20241029_083758_落石_80cm80cm_白天_晴_k411+530.mp4";
    // RTSP  = "/data/汉丹线数据-原始数据/创联/20241017/崩塌落石/20241017_104206_落石_50cm_白天_雨_k411+530.mp4";
    curandStatePhilox4_32_10_t* d_state;
    cudaStream_t stream;
   
    // cv::cuda::GpuMat frame_gpu;
    // cv::cuda::GpuMat gray_gpu(H, W, CV_8UC1);
    // cv::cuda::GpuMat gray_float_gpu(H, W, CV_32FC1);

    // 初始化 CUDA 流
    cudaStreamCreate(&stream);

    cudaMalloc((void**)&d_gray, H * W * sizeof(float));
    cudaMalloc((void**)&d_samples, H * W * num_samples * sizeof(float));
    cudaMalloc((void**)&d_img_count, H * W * sizeof(int));
    cudaMalloc((void**)&d_target_hit_count, H * W * sizeof(int));
    cudaMallocHost((void**)&h_mask_img, H * W * sizeof(int));
    cudaMalloc((void**)&d_mask_img, H * W * sizeof(int));
    cudaMalloc((void**)&d_state, H * W * sizeof(curandStatePhilox4_32_10_t));


    std::vector<int> h_offset = { -1, 0, 1, -1, 0, 1, -1, 0, 1 };
    thrust::device_vector<int> d_offset(h_offset.begin(), h_offset.end());

    dim3 thread_per_block(BLOSIZE, BLOSIZE);
    dim3 block_per_grid((W + BLOSIZE - 1) / BLOSIZE, (H + BLOSIZE - 1) / BLOSIZE);
    
    cv::VideoCapture Vs(RTSP);
    if (!Vs.isOpened()) {
        std::cout << "Error opening video stream or file" << std::endl;
        return -1;
    }
    cv::Mat frame,gray,gray_float,mask_img_8U_blur_binary,colored_image;
    int frame_id = 0;
    std::string savename = GetBaseName(RTSP);
    cv::VideoWriter writer(savename, cv::VideoWriter::fourcc('M','J','P','G'), 25, cv::Size(1920, 1080));
    
    while (true)
    {
        Vs >> frame;
        if (frame.empty()) break;

        // frame_gpu.upload(frame);
        cv::cvtColor(frame, gray, cv::COLOR_BGR2GRAY);
        gray.convertTo(gray_float, CV_32F); 
        cudaMemcpy(d_gray, gray_float.ptr<float>(), H * W * sizeof(float), cudaMemcpyHostToDevice);
        // 计算持续时间
        float milliseconds = 0;

        if (frame_id == 0) 
        {
            random_setup(d_state, time(NULL), H, W,BLOSIZE,stream);
            Init_Matarix(d_samples, d_img_count, d_target_hit_count, d_mask_img, num_samples,H,W,BLOSIZE,stream);
            ProcessFirstFrame(d_gray, d_samples, d_state, thrust::raw_pointer_cast(d_offset.data()), num_samples, H, W,BLOSIZE,stream);
        } 
        else 
        {
            // 开始计时
            cudaEvent_t start, end;
            cudaEventCreate(&start);
            cudaEventCreate(&end);
            cudaEventRecord(start, 0);

            ProcessNextFrame(d_gray, d_samples, d_img_count, d_target_hit_count, d_mask_img, d_state, thrust::raw_pointer_cast(d_offset.data()), H, W, min_match, random_sample, num_samples, max_hitcount,BLOSIZE,stream);
            cudaMemcpy(h_mask_img, d_mask_img, H * W * sizeof(int), cudaMemcpyDeviceToHost);
            
            // 记录结束事件
            cudaEventRecord(end, 0);
            // 等待事件完成
            cudaEventSynchronize(end);
            cudaEventElapsedTime(&milliseconds, start, end);
            // 打印结果
            std::cout<<"frame_id: "<<frame_id<< "   Elapsed time: " << " " << milliseconds << " milliseconds" << std::endl;

            // 释放事件
            cudaEventDestroy(start);
            cudaEventDestroy(end);

            // 将 h_mask_img 转换为 OpenCV Mat
            cv::Mat mask_img(H, W, CV_32S, h_mask_img);
            // 如果需要将数据转换为 CV_8U（例如，用于显示）
            cv::Mat mask_img_8U,mask_img_8U_blur;
            mask_img.convertTo(mask_img_8U, CV_8U);
            cv:medianBlur(mask_img_8U,mask_img_8U_blur,5);
            cv::threshold(mask_img_8U_blur, mask_img_8U_blur_binary,0,255, cv::THRESH_BINARY);
            fill_empty(mask_img_8U_blur_binary.clone(),mask_img_8U_blur_binary);
            colored_image = apply_mask_to_image(frame,mask_img_8U_blur_binary);
            // 将FPS转换为字符串
            std::string fpsText = "FPS: " + std::to_string(static_cast<int>(1000/milliseconds));
            // cv::putText(colored_image, fpsText, cv::Point(10, 30 ), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 255, 127), 3);
            // // // 等待按键
            // int key = cv::waitKey(10);
            // if (key == 'q') {
            //     break;  // 如果按下 'q' 键，退出循环
            // }
            // cv::namedWindow("VIBE_CUDA", cv::WINDOW_NORMAL);
            // cv::imshow("VIBE_CUDA",colored_image);
            writer.write(colored_image); // frame 是一个 cv::Mat 对象，代表视频帧
            Reset_Matarix(d_img_count, d_mask_img, H, W,BLOSIZE,stream);
        }
        cudaStreamSynchronize(stream);

        frame_id++;
        // std::cout<<"frame_id:"<<frame_id<<std::endl;
    }

    // 清理资源
    // cudaFree(d_gray);
    writer.release();
    cudaFree(d_samples);
    cudaFree(d_img_count);
    cudaFree(d_target_hit_count);
    cudaFree(d_state);
    cudaStreamDestroy(stream);
    return 0;
}
